package top.yudoge.dogeeditorserverj.llm.tools;


import dev.langchain4j.agent.tool.Tool;
import kotlin.Pair;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import top.yudoge.dogeeditorserverj.db.Column;
import top.yudoge.dogeeditorserverj.db.Table;
import top.yudoge.dogeeditorserverj.service.DataBaseMetaDataService;
import top.yudoge.dogeeditorserverj.service.TextSimilarityService;
import top.yudoge.dogeeditorserverj.transients.mro.TableColumnMRO;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

@Slf4j
@Component
public class DataBaseTools {

    @Autowired
    private DataBaseMetaDataService dataBaseMetaDataService;
    @Autowired
    private TextSimilarityService textSimilarityService;

    private final int K = 5;

    @Tool("Given multiple table name (and the datasource, database, schema currently in), it will return the columns of these table. If tableNames not empty, you should always use this first!")
    public List<TableColumnMRO> getColumns(
            String dataSourceName,
            String dataBaseName,
            String schemaName,
            List<String> tableNames
    ) {
        log.info("get Columns {}", tableNames);
        List<TableColumnMRO> mros = new ArrayList<>();
        for (String tableName : tableNames) {
            List<Column> columns = dataBaseMetaDataService.getColumns(dataSourceName, dataBaseName, schemaName, tableName);
            TableColumnMRO mro = new TableColumnMRO();
            mro.setTableName(tableName);
            mro.setColumns(columns);
            mros.add(mro);
        }
        return mros;
    }

    @Tool("Given some words that appear in user task that maybe a part of some table name or remark. " +
            "it will match tables use these words. And return the table and columns that matched!")
    public List<TableColumnMRO> getPossibleTable(
            String dataSourceName,
            String dataBaseName,
            String schemaName,
            List<String> words
    ) {
        log.info("get PossibleTable {}", words);

        List<Table> tables = dataBaseMetaDataService.getTable(dataSourceName, dataBaseName, schemaName);

        List<Table> similarTables = new ArrayList<>();

        for (String word : words) {

            PriorityQueue<Pair<Table, Double>> tableSimilarities = new PriorityQueue<>((a, b) -> {
                return a.component2().compareTo(b.component2());
            });

            for (Table table : tables) {
                String tableName = table.getTableName();
                tableSimilarities.add(new Pair<>(table, textSimilarityService.similarity(word, tableName)));
                String tableComment = table.getComment();
                tableSimilarities.add(new Pair<>(table, textSimilarityService.similarity(word, tableComment)));
                while (tableSimilarities.size() > 5) {
                    Pair<Table, Double> remove = tableSimilarities.remove();
                    log.debug("word [{}], {} to be evicted", word, remove);
                }
            }

            log.debug("word [{}], similar table list {}", tableSimilarities);

            Table table = null;
            while (!tableSimilarities.isEmpty()) {
                table = tableSimilarities.remove().component1();
            }

            if (table != null) {
                similarTables.add(table);
            }
        }

        return similarTables.stream().map(it -> {
            List<Column> columns = dataBaseMetaDataService.getColumns(dataSourceName, dataBaseName, schemaName, it.getTableName());
            TableColumnMRO mro = new TableColumnMRO();
            mro.setTableName(it.getTableName());
            mro.setColumns(columns);
            return mro;
        }).toList();
    }


}
